package com.atguigui.leetcode1;

import java.util.HashMap;
import java.util.Map;

/**
 * 2488.统计中位数为 K 的子数组
 * Project: leetcode
 * Package: com.atguigui.leetcode1
 * Version: 1.0
 * <p>
 * Created by WJX on 2023/3/16 8:49
 */
public class P2488CountSubarraysWithMedianK {

    public static void main(String[] args) {
        Solution solution = new P2488CountSubarraysWithMedianK().new Solution();
        // TO TEST
    }

    class Solution {
        public int countSubarrays(int[] nums, int k) {
            int ans = 0;
            int n = nums.length;
            int kIndex = -1;
            for (int i = 0; i < n; i++) {
                if (nums[i] == k) {
                    kIndex = i;
                    break;
                }
            }

            Map<Integer, Integer> counts = new HashMap<Integer, Integer>();
            counts.put(0, 1);
            int sum = 0;
            for (int i = 0; i < n; i++) {
                //转化数组
                sum += sign(nums[i] - k);
                if (i < kIndex) {
                    counts.put(sum, counts.getOrDefault(sum, 0) + 1);
                } else {
                    int prev0 = counts.getOrDefault(sum, 0);
                    int prev1 = counts.getOrDefault(sum - 1, 0);
                    ans += prev0 + prev1;
                }
            }


            return ans;
        }

        //将大于 kkk 的元素转换成 111，小于 kkk 的元素转换成 −1-1−1，等于 kkk 的元素转换成 000
        public int sign(int num) {
            if (num == 0) {
                return 0;
            }
            return num > 0 ? 1 : -1;
        }
    }
}
